﻿//#define USE_CODE_INJECTION

using System;
using System.Collections.Generic;
using System.Diagnostics;
using VMExplorer.MemoryObjects;
using System.Runtime.InteropServices;
using System.Timers;
using System.Threading;
using System.Windows.Forms;
using System.IO;
using VMExplorer.Win32;

namespace VMExplorer
{
	class VMEDocument
	{
		#region Private fields

		private Process _process;
		private SafeProcessHandle _hProcess;

		private List<IMemoryObject> _memObjects = new List<IMemoryObject> ();
		private List<RegionInfo> _addresses = new List<RegionInfo> ();

		private MemoryAddress _selectedAddress;
		private int _selectedPage;
		private int _selectedRegion;
		private IMemoryObject _selectedObject;

		private System.Timers.Timer autoUpdateTimer;

    private Scanner scanner;
    private VMStats stats;

    private IntPtr _nativeBufferForReadPage;

		#endregion

		#region Creation

		public VMEDocument ()
		{
      scanner = new Scanner (this);
      SectionObject.InstallScanRule (scanner);
      HeapObject.InstallScanRule (scanner);
      HeapSegmentObject.InstallScanRule (scanner);
      ModuleObject.InstallScanRule (scanner);
      ActxObject.InstallScanRule (scanner);

			_process = null;
			_hProcess = null;
			_selectedAddress = MemoryAddress.Invalid;
			_selectedPage = -1;
			_selectedRegion = -1;
			_selectedObject = null;

      _nativeBufferForReadPage = Marshal.AllocHGlobal (4096);

			autoUpdateTimer = new System.Timers.Timer (1000);
			autoUpdateTimer.Elapsed += delegate (object sender, System.Timers.ElapsedEventArgs e) { UpdateData (); };
		}

    public void Load (Process process)
    {
			_process = process;

			if (_hProcess!=null) _hProcess.Close ();
			_hProcess = Native.OpenProcess (
					Native.ProcessAccess.PROCESS_QUERY_INFORMATION
				| Native.ProcessAccess.PROCESS_VM_OPERATION
				| Native.ProcessAccess.PROCESS_VM_WRITE
				| Native.ProcessAccess.PROCESS_VM_READ
				| Native.ProcessAccess.PROCESS_CREATE_THREAD, false, (uint)_process.Id);

      UpdateData ();
      SelectedAddress = new MemoryAddress(0x00400);
    }

    ~VMEDocument ()
    {
      Marshal.FreeHGlobal (_nativeBufferForReadPage);
			if (_hProcess!=null) _hProcess.Close ();
    }

		#endregion

		#region Data update

		public void UpdateData ()
		{
			if(_hProcess.IsInvalid) return;

      lock (_addresses)
      {
        // Don't suspend current process !
        bool bSuspend = (_process.Id != Process.GetCurrentProcess ().Id);
			  if (bSuspend) Native.NtSuspendProcess (_hProcess);
			  

			  try
			  {
				  UpdateAddresses ();
				  UpdateMemoryObjects ();
				  LinkObjectsAndAddresses ();

          FindOrphans ();
				  LinkObjectsAndAddresses ();

				  ComputeVMStatistics ();
			  }
			  finally
			  {
				  if (bSuspend) Native.NtResumeProcess (_hProcess);
			  }
      }

			SelectedAddress = SelectedAddress;
			OnDataUpdated ();
		}

    private void FindOrphans ()
    {
      _memObjects.AddRange (scanner.FindOrphans ());
    }

    private void UpdateAddresses ()
    {
      _addresses.Clear ();
			if (_hProcess.IsInvalid) return;

      try
      {
        IntPtr pMem = IntPtr.Zero;
        Native.MEMORY_BASIC_INFORMATION mbi = new Native.MEMORY_BASIC_INFORMATION ();
        while (Native.VirtualQueryEx (_hProcess, pMem, out mbi, (UIntPtr)28) == /*sizeof(MEMORY_BASIC_INFORMATION)*/28)
        {
          pMem = (IntPtr)(pMem.ToInt32 () + (Int32)mbi.RegionSize);
          _addresses.Add (new RegionInfo (mbi));
        }
      }
      catch
      {
      }
    }

    private void UpdateMemoryObjects ()
		{
      if (_hProcess.IsInvalid) return;

			List<IMemoryObject> memObjects = new List<IMemoryObject> ();

			IntPtr unmanagedPtrForTBI = Marshal.AllocHGlobal (28);
      IntPtr unmanagedPtrForPBI = Marshal.AllocHGlobal (24);

			try
      {
        #region :: PEB ::

        try
        {
          Native.PROCESS_BASIC_INFORMATION pbi = new Native.PROCESS_BASIC_INFORMATION ();

          uint ntStatus = 0;
          uint returnLength = 0;
          ntStatus = Native.NtQueryInformationProcess (_hProcess, Native.PROCESSINFOCLASS.ProcessBasicInformation, unmanagedPtrForPBI, 24, ref returnLength);
          pbi = (Native.PROCESS_BASIC_INFORMATION)Marshal.PtrToStructure (unmanagedPtrForPBI, typeof (Native.PROCESS_BASIC_INFORMATION));
					if (ntStatus == 0)
					{
						memObjects.AddRange (PEBObject.Create (ProcessHandle, pbi.PebBaseAddress));
					}
				}
        catch
        {
        }

        #endregion
        #region :: Threads and stacks ::

        foreach (ProcessThread pt in _process.Threads)
				{
					IntPtr hThread = Native.OpenThread (Native.ThreadAccess.QUERY_INFORMATION, false, pt.Id);
					if (hThread != IntPtr.Zero)
					{
						try
						{
							Native.THREAD_BASIC_INFORMATION tbi = new Native.THREAD_BASIC_INFORMATION ();

							uint ntStatus = 0;
							uint returnLength = 0;
							ntStatus = Native.NtQueryInformationThread (hThread, Native.THREAD_INFORMATION_CLASS.ThreadBasicInformation, unmanagedPtrForTBI, 28, ref returnLength);
							tbi = (Native.THREAD_BASIC_INFORMATION)Marshal.PtrToStructure (unmanagedPtrForTBI, typeof (Native.THREAD_BASIC_INFORMATION));
							if (ntStatus == 0)
							{
								memObjects.AddRange (TEBObject.Create (ProcessHandle, tbi.TebBaseAddress, pt.Id));
							}
						}
						finally
						{
							Native.CloseHandle (hThread);
						}
					}
        }

        #endregion
      }
      finally
      {
        Marshal.FreeHGlobal (unmanagedPtrForTBI);
        Marshal.FreeHGlobal (unmanagedPtrForPBI);
      }

			lock (_memObjects)
			{
				_memObjects.Clear ();
				_memObjects.AddRange (memObjects);
			}

		}

		private void LinkObjectsAndAddresses ()
		{
      foreach (RegionInfo ri in _addresses)
      {
				if (ri.State == Native.PAGE_STATE.MEM_FREE)
					continue;

        foreach (IMemoryObject mo in _memObjects)
        {
					MemoryAddress baseAddres = mo.BaseAddress.LowerPageAligned;
					if (baseAddres == ri.Base || baseAddres == ri.Start)
          {
            if(!ri.Related.Contains(mo)) ri.Related.Add(mo);
            //break;
          }
        }

        if (ri.Related.Count>0)
        {
					foreach (IMemoryObject mo in ri.Related)
					{
						if (!mo.Ranges.Contains (ri)) mo.Ranges.Add (ri);
					}
        }
      }
    }

		private void ComputeVMStatistics ()
		{
      stats.Init ();
      foreach (RegionInfo ri in _addresses)
			{
        if (ri.State == Native.PAGE_STATE.MEM_FREE)
        {
          stats.sizeFree += ri.Size;
          continue;
        }

				switch (ri.State)
				{
          case Native.PAGE_STATE.MEM_RESERVE: stats.sizeReserved += ri.Size; break;
          case Native.PAGE_STATE.MEM_COMMIT: stats.sizeCommited += ri.Size; break;
				}

        switch (ri.Type)
        {
          case Native.PAGE_TYPE.MEM_IMAGE: stats.sizeImage += ri.Size; break;
          case Native.PAGE_TYPE.MEM_MAPPED: stats.sizeMapped += ri.Size; break;
          case Native.PAGE_TYPE.MEM_PRIVATE: stats.sizePrivate += ri.Size; break;
        }
			}
		}

		#endregion

		#region Public Properties

    public bool AutoUpdate
    {
      get { return autoUpdateTimer.Enabled; }
      set { autoUpdateTimer.Enabled = value; }
    }

    public List<IMemoryObject> MemoryObjects { get { return new List<IMemoryObject> (_memObjects); } }
		public List<RegionInfo> AddressRanges { get { return new List<RegionInfo> (_addresses); } }

		public int TotalPages 
    {
      get
      {
        int iTotalPages = 0;
        if (_addresses.Count > 0)
        {
          iTotalPages = _addresses [_addresses.Count - 1].Base.Page + (int)(_addresses [_addresses.Count - 1].Size / 4096);
        }
        return iTotalPages; 
      }
    }

		public MemoryAddress SelectedAddress
		{
			get { return _selectedAddress; }
			set
			{
				MemoryAddress oldAddress = _selectedAddress;
        int oldPage = oldAddress.Page;
				int oldRegion = _selectedRegion;
				IMemoryObject oldObject = _selectedObject;

        _selectedAddress = value.IsUserSpace ? value : MemoryAddress.Invalid;
				_selectedPage = _selectedAddress.Page;
        _selectedRegion = FindRegion (_selectedPage);
				_selectedObject = SelectedRange.Related.Count>0 ? SelectedRange.Related[0] : null;

				if (oldAddress != _selectedAddress)
				{
					OnSelectedAddressChanged (new SelectedAddressChangedEventArgs (oldAddress, _selectedAddress));
				}
				if (_selectedPage != oldPage)
				{
					OnPageSelectionChanged (new PageSelectionChangeEventArgs (oldPage, _selectedPage));
				}
				if (_selectedRegion != oldRegion)
				{
					OnRegionSelectionChanged (new PageSelectionChangeEventArgs (oldRegion, _selectedRegion));
				}
				if (_selectedObject != oldObject)
				{
					OnSelectedObjectChanged (new SelectedObjectChangeEventArgs (_selectedObject));
				}
			}
		}

		public int SelectedPage
		{
			get { return _selectedPage; }
			set { SelectedAddress = new MemoryAddress (value); }
		}

		public int SelectedRegion 
    {
      get { return _selectedRegion; } 
    }

		public RegionInfo SelectedRange
    {
      get { return (_selectedRegion != -1) ? _addresses [_selectedRegion] : RegionInfo.Empty; } 
    }

		public IMemoryObject SelectedObject 
    { 
      get { return _selectedObject; } 
    }

		public SafeProcessHandle ProcessHandle { get { return _hProcess; } }

		#endregion

		#region Events

		public event SelectedAddressChangedHandler SelectedAddressChanged;

		private void OnSelectedAddressChanged (SelectedAddressChangedEventArgs e)
		{
			if (SelectedAddressChanged != null)
			{
				SelectedAddressChanged (this, e);
			}
		}

		public event PageSelectionChangeHandler PageSelectionChanged;

		private void OnPageSelectionChanged (PageSelectionChangeEventArgs e)
		{
			if (PageSelectionChanged != null)
			{
				PageSelectionChanged (this, e);
			}
		}

		public event PageSelectionChangeHandler RegionSelectionChanged;

		private void OnRegionSelectionChanged (PageSelectionChangeEventArgs e)
		{
			if (RegionSelectionChanged != null)
			{
				RegionSelectionChanged (this, e);
			}
		}

		public event SelectedObjectChangeHandler SelectedObjectChanged;

		private void OnSelectedObjectChanged (SelectedObjectChangeEventArgs e)
		{
			if (SelectedObjectChanged != null)
			{
				SelectedObjectChanged (this, e);
			}
		}

		public event DataUpdatedHandler DataUpdated;

		private void OnDataUpdated ()
		{
			if (DataUpdated != null)
			{
				DataUpdated (this, new EventArgs ());
			}
		}

		#endregion

		#region Helpers

		public int FindRegion (int page)
		{
      if (page == -1) return -1;

      MemoryAddress target = new MemoryAddress (page);

      int select = 0;
			while ( select < _addresses.Count && !_addresses [select].Contains (target)) select++;
      if (select == _addresses.Count) select = -1;

			return select;
		}

    public bool ReadMemory (MemoryAddress baseAddress, IntPtr buffer, uint size)
		{
			return RemoteReader.ReadMemory (ProcessHandle, baseAddress, buffer, size);
		}


		// Assume size is less than 4096
    public ByteBuffer ReadByteBuffer (MemoryAddress baseAddress, int size)
    {
      // Check for invalid adress
      if (baseAddress == MemoryAddress.Invalid)
      {
        return null;
      }

			ByteBuffer retBuffer = new ByteBuffer ();
			retBuffer.BaseAddress = baseAddress;
			retBuffer.Size = size;

			int splitPoint = baseAddress.UpperPageOffset;
			int bytesToRead = 0;
			int readBytes = 0;

			// Read before split
			retBuffer.Start = 0;
			retBuffer.End = 0;
			bytesToRead = splitPoint;
			if (RemoteReader.ReadMemory (ProcessHandle, baseAddress, _nativeBufferForReadPage, (uint)bytesToRead))
			{
				retBuffer.End = bytesToRead;
				readBytes += bytesToRead;
			}
			else
			{
				retBuffer.Start = bytesToRead;
				retBuffer.End = bytesToRead;
			}

			// Read after split
			bytesToRead = size - splitPoint;
			if (RemoteReader.ReadMemory (ProcessHandle, baseAddress + splitPoint, (IntPtr)(_nativeBufferForReadPage.ToInt32 () + readBytes), (uint)bytesToRead))
			{
				retBuffer.End += bytesToRead;
				readBytes += bytesToRead;
			}

			// Convert unmanaged data to byte array
			retBuffer.Data = new byte [readBytes];
      unsafe
      {
				UnmanagedMemoryStream ums = new UnmanagedMemoryStream ((byte*)_nativeBufferForReadPage.ToPointer (), readBytes, readBytes, FileAccess.Read);
				ums.Read (retBuffer.Data, 0, readBytes);
      }

      return retBuffer;
    }

    #endregion
	}

  class SelectedAddressChangedEventArgs : EventArgs
  {
    private MemoryAddress _oldValue;
    private MemoryAddress _newValue;

    public SelectedAddressChangedEventArgs (MemoryAddress i_oldValue, MemoryAddress i_newValue)
    {
      _oldValue = i_oldValue;
      _newValue = i_newValue;
    }

    public MemoryAddress OldValue { get { return _oldValue; } }
    public MemoryAddress NewValue { get { return _newValue; } }
  };
  delegate void SelectedAddressChangedHandler (object sender, SelectedAddressChangedEventArgs e);

  class PageSelectionChangeEventArgs : EventArgs
  {
    private int _oldValue;
    private int _newValue;

    public PageSelectionChangeEventArgs (int i_oldValue, int i_newValue)
    {
      _oldValue = i_oldValue;
      _newValue = i_newValue;
    }

    public int OldValue { get { return _oldValue; } }
    public int NewValue { get { return _newValue; } }
  };
  delegate void PageSelectionChangeHandler (object sender, PageSelectionChangeEventArgs e);

  class SelectedObjectChangeEventArgs : EventArgs
  {
    private object _selection;

    public SelectedObjectChangeEventArgs (object i_selection)
    {
      _selection = i_selection;
    }

    public object Selection { get { return _selection; } }
  };
  delegate void SelectedObjectChangeHandler (object sender, SelectedObjectChangeEventArgs e);

  delegate void DataUpdatedHandler (object sender, EventArgs e);


  struct VMStats
  {
    public uint sizeFree;

    public uint sizeReserved;
    public uint sizeCommited;

    public uint sizeImage;
    public uint sizePrivate;
    public uint sizeMapped;

    public void Init ()
    {
      sizeFree = 0;

      sizeReserved = 0;
      sizeCommited = 0;

      sizeImage = 0;
      sizePrivate = 0;
      sizeMapped = 0;
    }
  };



}
